import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import LSTM, Dense, Dropout, Masking, Bidirectional
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.callbacks import ReduceLROnPlateau, EarlyStopping
from sklearn.utils.class_weight import compute_class_weight
from tensorflow.keras.preprocessing.sequence import pad_sequences
import json

# 数据集路径和生成模型的路径
DATASET_PATH = 'assets/database/main_database.json'  # 数据集路径
MODEL_SAVE_PATH = 'model/lstm_model.h5'  # 模型保存路径
KEY_TO_INDEX_PATH = 'model/lstm_key_to_index.json'  # key_to_index 保存路径
LABEL_TO_INDEX_PATH = 'model/lstm_label_to_index.json'  # label_to_index 保存路径

# 加载数据集
def load_data(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    return data

# 提取击键序列特征
def extract_sequence_features(data):
    sequences = []
    labels = []

    for entry in data:
        subject = entry['subject']
        keystrokes = entry['keystrokes']

        sequence = []
        for i in range(len(keystrokes)):
            key_content = keystrokes[i]['key_content']
            press_time = keystrokes[i]['press_time']
            release_time = keystrokes[i]['release_time']
            duration = release_time - press_time  # 按键持续时间
            if i > 0:
                interval = press_time - keystrokes[i - 1]['press_time']  # 击键间隔时间
                flight_time = press_time - keystrokes[i - 1]['release_time']  # 飞行时间
                prev_key_content = keystrokes[i - 1]['key_content']  # 前一个按键内容
                if interval == 0:
                    acceleration = 0.0
                else:
                    acceleration = (duration - (keystrokes[i - 1]['release_time'] - keystrokes[i - 1]['press_time'])) / interval  # 按键加速度
            else:
                interval = 0.0
                flight_time = 0.0
                acceleration = 0.0
                prev_key_content = 'None'  # 第一个按键没有前一个按键
            sequence.append([key_content, prev_key_content, press_time, release_time, duration, interval, flight_time, acceleration])

        sequences.append(sequence)
        labels.append(subject)

    return sequences, labels

# 将按键内容转换为整数编码
def encode_key_content(sequences):
    key_to_index = {}
    index = 0

    # 构建按键内容到整数的映射
    for sequence in sequences:
        for keystroke in sequence:
            key_content = keystroke[0]
            prev_key_content = keystroke[1]
            if key_content not in key_to_index:
                key_to_index[key_content] = index
                index += 1
            if prev_key_content not in key_to_index:
                key_to_index[prev_key_content] = index
                index += 1

    # 将按键内容转换为整数编码
    encoded_sequences = []
    for sequence in sequences:
        encoded_sequence = []
        for keystroke in sequence:
            key_content = keystroke[0]
            prev_key_content = keystroke[1]
            press_time = keystroke[2]
            release_time = keystroke[3]
            duration = keystroke[4]
            interval = keystroke[5]
            flight_time = keystroke[6]
            acceleration = keystroke[7]
            encoded_sequence.append([key_to_index[key_content], key_to_index[prev_key_content], press_time, release_time, duration, interval, flight_time, acceleration])
        encoded_sequences.append(encoded_sequence)

    return encoded_sequences, key_to_index

# 将标签转换为整数编码
def encode_labels(labels):
    label_to_index = {}
    index = 0

    # 构建标签到整数的映射
    for label in labels:
        if label not in label_to_index:
            label_to_index[label] = index
            index += 1

    # 将标签转换为整数编码
    encoded_labels = [label_to_index[label] for label in labels]

    return encoded_labels, label_to_index

# 主函数
def main():
    # 加载数据
    data = load_data(DATASET_PATH)

    # 提取击键序列特征
    sequences, labels = extract_sequence_features(data)

    # 将按键内容转换为整数编码
    encoded_sequences, key_to_index = encode_key_content(sequences)

    # 将标签转换为整数编码
    encoded_labels, label_to_index = encode_labels(labels)

    # 填充序列，使所有序列长度一致
    max_sequence_length = max(len(seq) for seq in encoded_sequences)  # 找到最长的序列长度
    X = pad_sequences(encoded_sequences, maxlen=max_sequence_length, padding='post', dtype='float32')  # 填充序列

    # 转换为 numpy 数组
    y = np.array(encoded_labels)

    # 将标签转换为 one-hot 编码
    y = to_categorical(y)

    # 分割数据集，确保每个类别的比例一致
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=np.argmax(y, axis=1))
    X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42, stratify=np.argmax(y_train, axis=1))

    # 计算类别权重
    class_weights = compute_class_weight('balanced', classes=np.unique(np.argmax(y_train, axis=1)), y=np.argmax(y_train, axis=1))
    class_weights = dict(enumerate(class_weights))

    # 构建 LSTM 模型
    model = Sequential()
    model.add(Masking(mask_value=0.0, input_shape=(X_train.shape[1], X_train.shape[2])))  # 忽略填充值
    model.add(Bidirectional(LSTM(256, return_sequences=True)))
    model.add(Dropout(0.5))
    model.add(Bidirectional(LSTM(128, return_sequences=True)))
    model.add(Dropout(0.5))
    model.add(Bidirectional(LSTM(64, return_sequences=False)))
    model.add(Dropout(0.5))
    model.add(Dense(128, activation='relu'))
    model.add(Dense(64, activation='relu'))
    model.add(Dense(y_train.shape[1], activation='softmax'))

    # 编译模型
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

    # 定义回调函数
    callbacks = [
        ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5, min_lr=0.00001),
        EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
    ]

    # 训练模型
    history = model.fit(X_train, y_train, epochs=50, batch_size=32, validation_data=(X_val, y_val), callbacks=callbacks, class_weight=class_weights)

    # 保存模型
    model.save(MODEL_SAVE_PATH)
    print(f"Model saved as '{MODEL_SAVE_PATH}'")

    # 保存 key_to_index
    with open(KEY_TO_INDEX_PATH, 'w') as f:
        json.dump(key_to_index, f)
    print(f"key_to_index saved as '{KEY_TO_INDEX_PATH}'")

    # 保存 label_to_index
    with open(LABEL_TO_INDEX_PATH, 'w') as f:
        json.dump(label_to_index, f)
    print(f"label_to_index saved as '{LABEL_TO_INDEX_PATH}'")

    # 评估模型
    y_pred = model.predict(X_test)
    y_pred = np.argmax(y_pred, axis=1)
    y_test_labels = np.argmax(y_test, axis=1)

    print("LSTM Accuracy:", accuracy_score(y_test_labels, y_pred))
    print("LSTM Classification Report:")
    print(classification_report(y_test_labels, y_pred))

if __name__ == "__main__":
    main()
    